import os
import random
import shutil


def split_files(path):
    # 获取指定路径下的所有文件名
    file_names = os.listdir(path)
    # 过滤掉文件夹，只保留文件
    file_names = [f for f in file_names if os.path.isfile(os.path.join(path, f))]
    # 打乱文件名列表，以保证随机性
    random.shuffle(file_names)
    total_files = len(file_names)
    # 计算训练集、验证集和测试集的文件数量
    train_count = int(total_files * 0.7)
    val_count = int(total_files * 0.15)
    test_count = total_files - train_count - val_count

    
    # 分割文件名列表
    train_files = file_names[:train_count]
    val_files = file_names[train_count:train_count + val_count]
    test_files = file_names[train_count + val_count:]

    # txtpath=label_path.replace('labels','trainvalsettxt')
    txtpath=path.replace('images','trainvalsettxt')
    if not os.path.exists(txtpath):  # 使用 os.path.exists() 检查路径是否存在
        os.makedirs(txtpath)  # 若路径不存在，则尝试创建多级目录
    # 写入训练集文件
    with open(os.path.join(txtpath,'train.txt'), 'w') as train_file:
        for file in train_files:
            train_file.write(file + '\n')

    # 写入验证集文件
    with open(os.path.join(txtpath,'val.txt'), 'w') as val_file:
        for file in val_files:
            val_file.write(file + '\n')

    # 写入测试集文件
    with open(os.path.join(txtpath,'test.txt'), 'w') as test_file:
        for file in test_files:
            test_file.write(file + '\n')

    return train_files, val_files, test_files


def move_files(path, train_files, val_files, test_files):
    # 创建 train、val、test 文件夹
    train_dir = os.path.join(path, 'train')
    val_dir = os.path.join(path, 'val')
    test_dir = os.path.join(path, 'test')
    os.makedirs(train_dir, exist_ok=True)
    os.makedirs(val_dir, exist_ok=True)
    os.makedirs(test_dir, exist_ok=True)

    # 移动训练集文件
    for file in train_files:
        src = os.path.join(path, file)
        dst = os.path.join(train_dir, file)
        shutil.move(src, dst)

    # 移动验证集文件
    for file in val_files:
        src = os.path.join(path, file)
        dst = os.path.join(val_dir, file)
        shutil.move(src, dst)

    # 移动测试集文件
    for file in test_files:
        src = os.path.join(path, file)
        dst = os.path.join(test_dir, file)
        shutil.move(src, dst)

def gen_images_set_and_move_file(path):
    train_files, val_files, test_files = split_files(path)
    move_files(path, train_files, val_files, test_files)
def move_files_based_on_txt(path,txt1,txt2,txt3):
    # 创建 train、val、test 文件夹
    train_dir = os.path.join(path, 'train')
    val_dir = os.path.join(path, 'val')
    test_dir = os.path.join(path, 'test')
    os.makedirs(train_dir, exist_ok=True)
    os.makedirs(val_dir, exist_ok=True)
    os.makedirs(test_dir, exist_ok=True)

    # 读取 train.txt 并移动文件
    with open(txt1, 'r') as f:
        train_files = [os.path.splitext(line.strip())[0] for line in f.readlines()]
        for file in os.listdir(path):
            file_name = os.path.splitext(file)[0]
            if file_name in train_files:
                src = os.path.join(path, file)
                dst = os.path.join(train_dir, file)
                shutil.move(src, dst)

    # 读取 val.txt 并移动文件
    with open(txt2, 'r') as f:
        val_files = [os.path.splitext(line.strip())[0] for line in f.readlines()]
        for file in os.listdir(path):
            file_name = os.path.splitext(file)[0]
            if file_name in val_files:
                src = os.path.join(path, file)
                dst = os.path.join(val_dir, file)
                shutil.move(src, dst)

    # 读取 test.txt 并移动文件
    with open(txt3, 'r') as f:
        test_files = [os.path.splitext(line.strip())[0] for line in f.readlines()]
        for file in os.listdir(path):
            file_name = os.path.splitext(file)[0]
            if file_name in test_files:
                src = os.path.join(path, file)
                dst = os.path.join(test_dir, file)
                shutil.move(src, dst)
if __name__ == "__main__":
    # 1. 首先注释下面的代码，只是gen_images_set_and_move_file
    # image_path = "data_zoo/splitsynbboxpcb/images"
    # image_path="data_zoo/pcb_truecolor_subcp4clsbbg640/images"
    image_path='/home/lhx/pcbsyn/data_zoo/pcb_truecolor_subcpjrsbg/images'
    # gen_images_set_and_move_file(image_path)


    # 2. 然后选择下面的特定path，运行move_files_based_on_txt
    # label_path="data_zoo/splitsynbboxpcb/labels"
    # txt1="data_zoo/splitsynbboxpcb/trainvalsettxt/train.txt"
    # txt2="data_zoo/splitsynbboxpcb/trainvalsettxt/val.txt"
    # txt3="data_zoo/splitsynbboxpcb/trainvalsettxt/test.txt"


    # image_path = "data_zoo/pcbsubcp4clsbbg640/images"
    


    # label_path="data_zoo/pcbsubcp4clsbbg640/labels"
    # txt1="data_zoo/pcbsubcp4clsbbg640/trainvalsettxt/train.txt"
    # txt2="data_zoo/pcbsubcp4clsbbg640/trainvalsettxt/val.txt"
    # txt3="data_zoo/pcbsubcp4clsbbg640/trainvalsettxt/test.txt"


    # image_path = "data_zoo/pcb_truecolor_subcp4clsbbg640/images"
    # label_path="data_zoo/pcb_truecolor_subcp4clsbbg640/labels"
    # bbox_path="data_zoo/pcb_truecolor_subcp4clsbbg640/bboxs"
    # mask_path="data_zoo/pcb_truecolor_subcp4clsbbg640/mask"
    # txt1="data_zoo/pcb_truecolor_subcp4clsbbg640/trainvalsettxt/train.txt"
    # txt2="data_zoo/pcb_truecolor_subcp4clsbbg640/trainvalsettxt/val.txt"
    # txt3="data_zoo/pcb_truecolor_subcp4clsbbg640/trainvalsettxt/test.txt"
    # image_path = "data_zoo/splitsynpcbwithjrsdepbg/images"
    # label_path="data_zoo/splitsynpcbwithjrsdepbg/labels"
    # bbox_path="data_zoo/splitsynpcbwithjrsdepbg/bboxs"
    # mask_path="data_zoo/splitsynpcbwithjrsdepbg/mask"
    # txt1="data_zoo/splitsynpcbwithjrsdepbg/trainvalsettxt/train.txt"
    # txt2="data_zoo/splitsynpcbwithjrsdepbg/trainvalsettxt/val.txt"
    # txt3="data_zoo/splitsynpcbwithjrsdepbg/trainvalsettxt/test.txt"
    # image_path = "data_zoo/pcb_truecolor_subcpjrsbg/images"
    image_path="data_zoo/pcb_truecolor_subcpjrsbg/depths"
    label_path="data_zoo/pcb_truecolor_subcpjrsbg/depths"
    bbox_path="data_zoo/pcb_truecolor_subcpjrsbg/bboxs"
    mask_path="data_zoo/pcb_truecolor_subcpjrsbg/mask"
    txt1="data_zoo/pcb_truecolor_subcpjrsbg/trainvalsettxt/train.txt"
    txt2="data_zoo/pcb_truecolor_subcpjrsbg/trainvalsettxt/val.txt"
    txt3="data_zoo/pcb_truecolor_subcpjrsbg/trainvalsettxt/test.txt"


    txtpath=label_path.replace('labels','trainvalsettxt')
    if not os.path.exists(txtpath):  # 使用 os.path.exists() 检查路径是否存在
        os.makedirs(txtpath)  # 若路径不存在，则尝试创建多级目录
    # gen_images_set_and_move_file(image_path)
    move_files_based_on_txt(label_path,txt1,txt2,txt3)


    #TLDR
    """
    先执行gen_images_set_and_move_file(image_path)注释move_files_based_on_txt(label_path,txt1,txt2,txt3)
    最后注释gen_images_set_and_move_file执行move_files_based_on_txt
    """
    

    